/*
 * Ex:			
 * File name:   scheduler.cpp
 *
 *  Created on: Apr 3, 2012
 *      Author: ayash
 *
 *
 * Descriptor:
 *
 */

#include <setjmp.h>
#include <signal.h>
#include <cassert>
#include <iostream>
#include <cstring>

#include "common_defs.h"
#include "user_thread.h"
#include "scheduler.h"
#include "uthreads.h"


using namespace std;

// this instance will be used as comparator of threads' pointers
ThreadsComparator threads_comarator;


////////////////////////////////////////////////////////////////////////
#ifdef __x86_64__
// code for 64 bit Intel arch

typedef unsigned long address_t;
#define JB_SP 6
#define JB_PC 7

//A translation required when using an address of a variable
//Use this as a black box in your code.
address_t translate_address(address_t addr)
{
    address_t ret;
    asm volatile("xor    %%fs:0x30,%0\n"
		"rol    $0x11,%0\n"
                 : "=g" (ret)
                 : "0" (addr));
    return ret;
}

#else
// code for 32 bit Intel arch

typedef unsigned int address_t;
#define JB_SP 4
#define JB_PC 5

//A translation required when using an address of a variable
//Use this as a black box in your code.
address_t translate_address(address_t addr)
{
    address_t ret;
    asm volatile("xor    %%gs:0x18,%0\n"
		"rol    $0x9,%0\n"
                 : "=g" (ret)
                 : "0" (addr));
    return ret;
}

#endif

////////////////////////////////////////////////////////////////////////

// disable previous virtual timer setups
static result_t disable_timer()
{
	struct itimerval tv;
	tv.it_value.tv_sec = 0;
	tv.it_value.tv_usec = 0;
	tv.it_interval.tv_sec = 0;
	tv.it_interval.tv_usec = 0;

	int ret_code = setitimer(ITIMER_VIRTUAL, & tv, NULL);
	if ( 0 != ret_code )
	{
		cerr<<"system error: disabling timer of a thread" << endl;
		return ERROR_SYSTEM_CALL;
	}

	return OPERATION_SUCCESS;
}

// perform a virtual timer signal request after num_millisecs milliseconds
static result_t enable_timer( uint num_millisecs )
{
	itimerval 	tv;
	uint		ret_code;

	// Set timer struct values to the requested time.
    tv.it_value.tv_sec = num_millisecs / 1000; // Convert to secs.
    tv.it_value.tv_usec = (num_millisecs % 1000) * 1000;
    tv.it_interval.tv_sec = 0;
    tv.it_interval.tv_usec = 0;

    ret_code = setitimer( ITIMER_VIRTUAL, & tv, NULL );
    if( 0 != ret_code )
    {
    	cerr<<"system error: setting timer for thread" << endl;
    	return ERROR_SYSTEM_CALL;
    }

    return OPERATION_SUCCESS;

}

Scheduler:: Scheduler():
		_threads_counter(0), _running(NULL), _sleeping (NULL), _deleted_thread(NULL)
{
}

Scheduler:: ~Scheduler()
{
	map<uint,UserThread*>:: iterator it;

	// delete all the threads' instances that weren't terminated already
	for( it = _all_threads.begin(); it != _all_threads.end(); it++ )
	{
		assert( NULL != it->second );
		delete it->second;
	}

	if ( NULL != _deleted_thread )
	{
		delete _deleted_thread;
	}


}

inline result_t Scheduler:: push_to_ready_heap( UserThread *thread )
{
	result_t result = OPERATION_SUCCESS;

	assert( NULL != thread );

	DBG_PRINT("Pushing thread ID: %d and State: %d \n",
			  thread->get_thread_id(),
			  thread->get_thread_state());

	/*
	 * this check allows pushing a thread that was popped from the ready heap
	 * and a system call failed before it was scheduled - we should push the
	 * given thread back to the heap without changing it's ready time
	 */
	if ( THREAD_STATE_READY != thread->get_thread_state() )
	{
		result = thread->set_thread_state( THREAD_STATE_READY );
	}

	_ready.push_back(thread);

	push_heap( _ready.begin(),_ready.end() , threads_comarator);

	return result;
}

void Scheduler:: perform_context_switch( UserThread *prev_running,
										 UserThread *to_schedule )
{
	assert( NULL != prev_running && NULL != to_schedule );

	DBG_PRINT("prev: id = %u, priority = %u\t new: id = %u, priority = %u\n",
			  prev_running->get_thread_id(),
			  prev_running->get_thread_priority(),
			  to_schedule->get_thread_id(),
			  to_schedule->get_thread_priority());


	int ret_val = sigsetjmp( *(prev_running->get_thread_jmp_buffer()) , 1);
	if ( 1 == ret_val )
	{
		if ( NULL != _deleted_thread )
		{
			delete _deleted_thread;
			_deleted_thread = NULL;
		}

		return;
	}

	siglongjmp( *(to_schedule->get_thread_jmp_buffer()) , 1);
}

/*
 * check if the given thread should switch with the running
 * in case the given thread has higher priority, the running thread will be enter
 * to the ready heap and the given thread will be marked as running (context switch will
 * be made too) otherwise the given thread will be added to the ready heap
 */
result_t Scheduler:: check_vs_running( UserThread *thread )
{
	assert( NULL != thread );

	result_t result;
	bool perfor_context_switch = false;
	UserThread *thread_to_heap = NULL;

	DBG_PRINT("running: id = %u, priority = %u\t new: id = %u, priority = %u\n",
			  _running->get_thread_id(),
			  _running->get_thread_priority(),
			  thread->get_thread_id(),
			  thread->get_thread_priority());

	if ( ! threads_comarator.operator ()( thread, _running ) )
	{
		/*
		 * the given thread is with higher priority than the running -
		 * perform a context switch
		 *
		 */
		thread_to_heap = _running;

		perfor_context_switch = true;

	}
	else
	{
		thread_to_heap = thread;
	}

	result = push_to_ready_heap( thread_to_heap );
	if( OPERATION_SUCCESS != result )
	{
		return result;
	}

	if ( perfor_context_switch )
	{
		_running = thread;
		_running -> set_thread_state( THREAD_STATE_RUNNING );

		perform_context_switch( thread_to_heap, _running);
	}

	return OPERATION_SUCCESS;
}

/*
 *  keep the context of the main thread - this method should be called once
 *  during initialization in order to add the main thread to the scheduler database
 */
void Scheduler:: keep_main_thread_u_area()
{
	UserThread *main_thread;
	sigjmp_buf main_jbuf;

	// keep main's u_area
	sigsetjmp( main_jbuf, 1 );
	

	main_thread = new UserThread( _threads_counter++ , MAX_PRI / 2 );
	main_thread->set_thread_u_area( & main_jbuf );
	assert( ! main_thread->get_thread_id() ); // main thread must be with id == 0

	_all_threads.insert(  pair<uint,UserThread*>( 0,main_thread) );
	_running = main_thread;
	_running->set_thread_state( THREAD_STATE_RUNNING );

}

uint Scheduler:: get_running_thread_id() const
{
	assert( NULL != _running );


	return _running->get_thread_id();
}

UserThread* Scheduler:: pop_top_heap ()
{
	assert( ! _ready.empty() );

	UserThread *top_thread = _ready.front();

	pop_heap ( _ready.begin(), _ready.end(), threads_comarator );

	_ready.pop_back();

	return top_thread;

}

/*
 *  perform a context switch between the top of the ready heap and the current running
 *  thread.
 */
result_t Scheduler:: schedule_from_heap ()
{
	result_t result;

	assert( ! _ready.empty() );

	UserThread *top_thread = pop_top_heap();

	result = push_to_ready_heap(_running);
	if ( OPERATION_SUCCESS != result )
	{
		return result;
	}

	UserThread *running_thread = _running;

	_running = top_thread;
	_running->set_thread_state( THREAD_STATE_RUNNING );

	perform_context_switch( running_thread, _running );

	return result;
}

result_t Scheduler:: set_thread_priority ( uint tid, uint new_priority )
{
	UserThread *top_thread = NULL;
	uint prev_priority;

	// scheduling decision should be made only if the thread is running or ready
	if ( NULL != _sleeping && tid == _sleeping->get_thread_id() )
	{
		_sleeping->set_thread_priority( new_priority );

		return OPERATION_SUCCESS;
	}

	if ( tid == _running->get_thread_id() )
	{
		prev_priority = _running->set_thread_priority( new_priority );

		if ( prev_priority < new_priority )
		{
			/*
			 * the new priority is lesser than the old one - there might be thread
			 * in the ready heap that its' priority is higher than the current running
			 * thread
			 */
			top_thread = _ready.front();

			if ( ! threads_comarator.operator ()( top_thread, _running )  )
			{
				/*
				 *  the priority of the thread at the top of the ready heap is higher
				 *  switch between the running and the top of the heap
				 */
				return schedule_from_heap();
			}
		}


		return OPERATION_SUCCESS;
	}

	/*
	 * the thread with the given ID is not running or sleeping - find it in the list of
	 * all threads
	 */
	UserThread *thread = get_thread_by_id( tid );

	if ( NULL == thread )
	{
		return ERROR_THREAD_UNKNOWN;
	}

	// update the priority
	prev_priority = thread->set_thread_priority( new_priority );

	thread_state_t thread_state = thread->get_thread_state();
	if ( THREAD_STATE_SUSPENDED == thread_state )
	{
		return OPERATION_SUCCESS;
	}

	/*
	 * the thread with the given ID is in the ready heap - check if a context switch
	 * is needed
	 */
	assert( THREAD_STATE_READY == thread_state );

	/*
	 * first we have to rearrange the heap so it be valid after we changed the priority
	 * we can't remove the thread from the ready heap before we rearrange it - even
	 * if we know that the thread going to switch the current running thread
	 */
	make_heap( _ready.begin(), _ready.end(), threads_comarator );

	top_thread = _ready.front();

	if ( ! threads_comarator.operator ()( top_thread, _running )  )
	{
		/*
		 *  the priority of the thread at the top of the ready heap is higher -
		 *  switch between the running and the top of the heap
		 */

		/*
		 * only the thread that we changed its' priority can have higher priority than
		 * the thread that is currently running
		 */
		assert( tid == top_thread->get_thread_id() );

		return schedule_from_heap();
	}

	return OPERATION_SUCCESS;
}

result_t Scheduler:: thread_resume (uint tid)
{
	result_t result;

	// scheduling decision should be made only if the thread is suspended or sleeping
	if ( _running->get_thread_id() == tid )
	{
		// nothing to do
		return OPERATION_SUCCESS;
	}

	if ( NULL != _sleeping && tid == _sleeping->get_thread_id() )
	{
		/*
		 * disable the last configuration of the timer so the resumed thread won't become
		 * ready by it's own.
		 */
		result = disable_timer();
		if ( OPERATION_SUCCESS != result )
		{
			return result;
		}

		UserThread *thread = _sleeping;

		_sleeping = NULL;

		return check_vs_running( thread );

	}

	/*
	 * the thread with the given ID is not running or sleeping - find it in the list of
	 * all threads
	 */
	UserThread *thread = get_thread_by_id( tid );
	if ( NULL == thread )
	{
		return ERROR_THREAD_UNKNOWN;
	}

	if ( THREAD_STATE_READY == thread->get_thread_state() )
	{
		// nothing to do
		return OPERATION_SUCCESS;
	}

	assert( THREAD_STATE_SUSPENDED == thread->get_thread_state() );

	return check_vs_running( thread );
}

result_t Scheduler:: thread_create (void (*thread_func)(void), uint priority, int *id)
{
	result_t		ret_code;
	int				jmp_code;
	UserThread		*new_thread;
	sigjmp_buf		jbuf;
	address_t		sp,pc;

	assert( NULL != thread_func && NULL != id );

	// keep the ID of the created thread in the user pointer
	*id = _threads_counter++;

	jmp_code = sigsetjmp( jbuf, 1);

	if(1 == jmp_code)
	{
		// This code is executed if we are returning to this after a long jump.
		// We need to discuss this. Why should we get here ? I think it is a problem
		// if we get here as we are changing the SP and PC
		ERR_PRINT("RETURNING FROM JMP TO SPAWN %d",*id);
	}

    new_thread = new UserThread( *id, priority);
    sp = (address_t)new_thread->get_stack_addr() + STACK_SIZE - sizeof(address_t);
    pc = (address_t)thread_func;
    (jbuf[0].__jmpbuf)[JB_SP] = translate_address(sp);
    (jbuf[0].__jmpbuf)[JB_PC] = translate_address(pc);

        if ( sigemptyset( & jbuf[0].__saved_mask )
        		// handle corner case when a new thread schedule when the signals are off
        		|| sigdelset( & jbuf[0].__saved_mask, SIGVTALRM ) )
    {
    	delete new_thread;
    	cerr <<"system error: failed to set empty mask" << endl;
    	return ERROR_SYSTEM_CALL;
    }

    new_thread->set_thread_u_area( & jbuf );

	/* Update relevant structures */
	_all_threads.insert(  pair<uint,UserThread*>( *id,new_thread) );
	ret_code = check_vs_running( new_thread );
	if( OPERATION_SUCCESS != ret_code )
	{
		_all_threads.erase( *id );
		*id = -1; // mark the id as invalid
		_threads_counter--;
		delete new_thread;

		return ret_code;
	}

	return OPERATION_SUCCESS;
}

result_t Scheduler:: thread_suspend (uint tid)
{
	UserThread*		suspended_thread;
	result_t		res_code;

	if ( tid == _running->get_thread_id() )
	{
		_running->set_thread_state( THREAD_STATE_SUSPENDED );

		suspended_thread = _running;

		UserThread* top_heap = pop_top_heap();

		_running = top_heap;
		_running->set_thread_state(THREAD_STATE_RUNNING);

		perform_context_switch( suspended_thread ,_running);

		return OPERATION_SUCCESS;
	}

	if ( NULL !=  _sleeping && tid == _sleeping->get_thread_id() )
	{
		res_code = disable_timer();
		if ( OPERATION_SUCCESS != res_code )
		{
			return res_code;
		}

		_sleeping->set_thread_state( THREAD_STATE_SUSPENDED );

		_sleeping = NULL;

		return OPERATION_SUCCESS;

	}

	suspended_thread = get_thread_by_id(tid);
	if( NULL == suspended_thread )
	{
		return ERROR_THREAD_UNKNOWN;
	}

	if ( THREAD_STATE_SUSPENDED == suspended_thread->get_thread_state() )
	{
		// nothing to do
		return OPERATION_SUCCESS;
	}

	assert( THREAD_STATE_READY == suspended_thread->get_thread_state() );

	// Thread must be in ready data structure
	remove_from_ready( suspended_thread );

	suspended_thread->set_thread_state( THREAD_STATE_SUSPENDED );

	return OPERATION_SUCCESS;
}

result_t Scheduler:: thread_sleep (uint tid, uint num_millisecs)
{
	UserThread* top_thread = NULL;
    result_t	result = OPERATION_SUCCESS;

	if( NULL != _sleeping )
	{
		// A Thread is already sleeping. Only one is allowed
		ERR_PRINT("the thread with ID #%u already sleeping\n", _sleeping->get_thread_id());
		return ERROR_THREAD_ALREADY_SLEEPING;
	}

	// check if the given thread is the current that running
	if ( tid == _running->get_thread_id() )
	{
		top_thread = pop_top_heap();

		_sleeping = _running;

	}
	else
	{
		_sleeping = get_thread_by_id( tid );

		if( NULL == _sleeping )
		{
			return ERROR_THREAD_UNKNOWN;
		}

		switch ( _sleeping->get_thread_state() ) {

		case THREAD_STATE_READY:

			remove_from_ready( _sleeping );

			break;

		case THREAD_STATE_SUSPENDED:
			// nothing to do here - we only need to put the thread to sleep
			break;

		default:
			assert(0);
			break;

		}
	}

    result = enable_timer(num_millisecs);
    if ( OPERATION_SUCCESS != result )
    {
    	if ( NULL != top_thread )
    	{
    		push_to_ready_heap(top_thread);

    	}
    	else if ( THREAD_STATE_READY == _sleeping->get_thread_state() )
    	{
    		push_to_ready_heap( _sleeping );
    	}

    	_sleeping = NULL;

    	return ERROR_SYSTEM_CALL;
    }

    _sleeping->set_thread_state( THREAD_STATE_SLEEP );

    if ( NULL != top_thread )
    {
    	_running = top_thread;

    	_running->set_thread_state( THREAD_STATE_RUNNING );

    	perform_context_switch( _sleeping , _running );
    }

	return result;
}

result_t Scheduler:: thread_terminate (uint tid)
{
	UserThread* top_thread;
	UserThread* removed_thread;
	result_t	res_code = OPERATION_SUCCESS;
	uint num_erased = 0; // used for validation check


	(void) num_erased; // make the compiler happy
	assert( tid );

	// Remove from the general thread map
	removed_thread = get_thread_by_id( tid );
	if ( NULL == removed_thread )
	{
		return ERROR_THREAD_UNKNOWN;
	}

	num_erased = _all_threads.erase( tid );
	assert( num_erased );

	switch( removed_thread->get_thread_state() ) {

	case THREAD_STATE_RUNNING:

		if ( NULL != _deleted_thread )
		{
			delete _deleted_thread;
		}

		_deleted_thread = _running;

		top_thread = pop_top_heap();

		_running = top_thread;

		_running->set_thread_state( THREAD_STATE_RUNNING );

		siglongjmp( *(_running->get_thread_jmp_buffer()), 1 );

		// shouldn't never be here
		break;

	case THREAD_STATE_READY:

		remove_from_ready(removed_thread);

		break;

	case THREAD_STATE_SUSPENDED:

		// Nothing to do here - we only need to delete the thread database

		break;

	case THREAD_STATE_SLEEP:

		res_code = disable_timer();

        _sleeping = NULL;

		break;

	default:

		assert(0); // This can't happen. What should we do in such cases
		break;
	}

	delete removed_thread;

	return res_code;

}

UserThread* Scheduler::get_thread_by_id(uint tid)
{
	map<uint,UserThread*>:: iterator it = _all_threads.find( tid );

	if ( _all_threads.end() == it )
	{
		return NULL;
	}

	return it->second;
}

void Scheduler:: remove_from_ready(UserThread* thread)
{
	UserThread *poped_thread; // will be used for validation check

	(void)poped_thread; //make the compiler happy

	assert( NULL != thread && THREAD_STATE_READY == thread->get_thread_state() );

	uint prev_priority = thread->set_thread_priority( MIN_VAL_PRIORITY );

	/*
	 * rearrange the heap after we updated the given thread priority so it be at the
	 * top of the heap - this allows up to pop it from the heap
	 */
	make_heap( _ready.begin(), _ready.end(), threads_comarator );

	poped_thread = pop_top_heap();

	assert( poped_thread->get_thread_id() == thread->get_thread_id() );

	// restore the previous priority of the thread after we removed it from the ready heap
	thread->set_thread_priority( prev_priority );
}

void Scheduler:: vtimer_handler()
{
	result_t	ret_code;
	UserThread *sleeping_thread;

	assert( NULL != _sleeping );

	sleeping_thread = _sleeping;
	_sleeping = NULL;
	ret_code = check_vs_running( sleeping_thread );

	if(OPERATION_SUCCESS != ret_code)
	{
		ERR_PRINT("Cannot resume thread in handler due to %d",ret_code);
	}
}
